import uvicorn
from fastapi import FastAPI, Body

from milvus.milvus_op import MilvusOP

milvus_op = MilvusOP(db_name='state_vector_db', collection_name="hybrid")
app = FastAPI()

config = {'ns': 1, 'ts': 1, 'e': 1, 'td': 1}


@app.post("/queries")
def retrieval_milvus(texts=Body(..., description='字符串列表')):
    queries = texts['queries']
    res = milvus_op.hybrid_search(queries, limit=3)
    res = [
        {'metadata':
             {'_source': 'knowledge', 'segment_id': i['id'], "retriever_from": "workflow", 'score': i['score'],
              'mark': i['remark']},
         'title': i['name'] + i['type'],
         'content': i['text']
         }
        for i in res
    ]
    return {'result': res}


uvicorn.run(app, host='172.16.12.239', port=7779)
